# @Version        : 1.0
# @Update Time    : 2024/12/9 22:25
# @File           : log_middleware.py
# @IDE            : PyCharm
# @Desc           : 日志中间件，用于记录请求和响应的详细信息

import time
import json
from typing import Iterable, Optional
from urllib.parse import unquote

from starlette.types import ASGIApp, Receive, Scope, Send
from applications.log import logger


class LogMiddleware:
    """
    访问日志中间件（ASGI）
    - 统一接管请求/响应日志
    - 自动截断请求体、脱敏敏感字段
    - 分级输出：2xx/3xx -> info，4xx -> warning，5xx -> error
    - 统计响应字节数，记录耗时
    - 不依赖私有 API，兼容流式/大包场景
    """

    def __init__(
        self,
        app: ASGIApp,
        *,
        max_body: int = 4096,
        log_request_body: bool = True,
        include_query: bool = True,
        skip_paths: Optional[Iterable[str]] = ("/docs", "/openapi.json", "/redoc", "/favicon.ico"),
        skip_content_types: Optional[Iterable[str]] = ("multipart/form-data", "application/octet-stream"),
    ):
        self.app = app
        self.max_body = max_body
        self.log_request_body = log_request_body
        self.include_query = include_query
        self.skip_paths = set(skip_paths or [])
        self.skip_content_types = tuple(skip_content_types or ())

    async def __call__(self, scope: Scope, receive: Receive, send: Send) -> None:
        if scope["type"] != "http" or scope["path"] in self.skip_paths:
            await self.app(scope, receive, send)
            return

        start = time.perf_counter()
        req_body_sample = b""
        resp_bytes = 0
        status_code = 200

        headers = {k.decode().lower(): v.decode() for k, v in scope.get("headers", [])}
        content_type = headers.get("content-type", "")
        method = scope["method"]
        path = scope["path"]
        query_string = scope.get("query_string", b"").decode()
        client = scope.get("client")
        client_ip = client[0] if client else "-"
        ua = headers.get("user-agent", "-")
        http_version = f"http/{scope.get('http_version', '1.1')}"

        async def receive_wrapper():
            nonlocal req_body_sample
            message = await receive()
            if message["type"] == "http.request":
                body = message.get("body", b"")
                if body and len(req_body_sample) < self.max_body:
                    need = self.max_body - len(req_body_sample)
                    req_body_sample += body[:need]
            return message

        async def send_wrapper(message):
            nonlocal resp_bytes, status_code
            if message["type"] == "http.response.start":
                status_code = message["status"]
            elif message["type"] == "http.response.body":
                body = message.get("body", b"") or b""
                resp_bytes += len(body)
            await send(message)

        try:
            await self.app(scope, receive_wrapper, send_wrapper)
        finally:
            dur_ms = int((time.perf_counter() - start) * 1000)

            # 解析请求体（按需，JSON 且不过大）
            body_obj = None
            if self.log_request_body and "application/json" in content_type and req_body_sample:
                body_obj = self._try_parse_json(req_body_sample)
                body_obj = self._sanitize(body_obj)

            # 组织结构化日志
            payload = {
                "method": method,
                "path": path,
                "status": status_code,
                "http": http_version,
                "client_ip": client_ip,
                "ua": ua,
                "duration_ms": dur_ms,
                "content_length": resp_bytes or None,
                "content_type": content_type or None,
            }
            if self.include_query and query_string:
                payload["query"] = unquote(query_string)
            if body_obj is not None:
                payload["request_body"] = body_obj

            # 分级输出
            log = logger.bind(biz="access", **payload)
            if status_code >= 500:
                log.error(f"{method} {path} -> {status_code} in {dur_ms}ms")
            elif status_code >= 400:
                log.warning(f"{method} {path} -> {status_code} in {dur_ms}ms")
            else:
                log.info(f"{method} {path} -> {status_code} in {dur_ms}ms")

    @staticmethod
    def _try_parse_json(data: bytes):
        try:
            return json.loads(data)
        except Exception:
            try:
                return data.decode("utf-8", errors="replace")
            except Exception:
                return "<unparseable>"

    @staticmethod
    def _sanitize(obj):
        # 递归脱敏敏感字段
        sensitive = {"password", "passwd", "pwd", "token", "access_token", "authorization",
                     "cookie", "set-cookie", "secret", "api_key", "client_secret"}
        def mask(val):
            s = str(val)
            if len(s) <= 6:
                return "***"
            return s[:3] + "***" + s[-2:]

        def rec(x):
            if isinstance(x, dict):
                return {k: ("***" if k.lower() in sensitive else rec(v)) for k, v in x.items()}
            if isinstance(x, list):
                return [rec(i) for i in x]
            if isinstance(x, (str, int, float, bool)) or x is None:
                return x
            return str(x)
        return rec(obj)
